package org.jeecg.modules.airag.llm.splitter;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.data.segment.TextSegment;
import org.jeecg.common.util.oConvertUtils;

import java.util.ArrayList;
import java.util.List;
import java.util.regex.Pattern;

/**
 * 功能描述
 *
 * @program: jeecg-boot
 * @description:
 * @author: wwq
 * @create: 2025-08-11 10:23
 **/
public class CustomDelimiterSplitter {


    /**
     * 创建自定义分隔符分段器
     *
     * @param delimiter 分隔符
     * @param maxSegmentSize 最大分段大小
     * @param maxOverlapSize 最大重叠大小
     * @return 文档分段器
     */
    public static DocumentSplitter FixedCharSplitter(String delimiter, int maxSegmentSize, int maxOverlapSize) {
        return new DocumentSplitter() {
            @Override
            public List<TextSegment> split(Document document) {
                String text = document.text();
                List<TextSegment> segments = new ArrayList<>();

                // 按分隔符分割文本
                String[] parts = text.split(Pattern.quote(delimiter));

                for (int i = 0; i < parts.length; i++) {
                    String part = parts[i].trim();
                    if (oConvertUtils.isEmpty(part)) {
                        continue;
                    }

                    // 如果分段太长，进一步分割
                    if (part.length() > maxSegmentSize) {
                        List<TextSegment> subSegments = splitLongText(part, maxSegmentSize, maxOverlapSize, document.metadata());
                        segments.addAll(subSegments);
                    } else {
                        // 创建元数据副本并添加索引
                        Metadata segmentMetadata = document.metadata().copy();
                        segmentMetadata.put("index", i);
                        segments.add(TextSegment.from(part, segmentMetadata));
                    }
                }

                return segments;
            }
        };
    }

    /**
     * 创建父子分段器
     * 父分段用于提供完整上下文，子分段用于精确检索
     *
     * @param parentDelimiter 父分隔符，如"####"
     * @param childDelimiter 子分隔符，如"@@@"
     * @param maxSegmentSize 最大分段大小
     * @param maxOverlapSize 最大重叠大小
     * @return 父子分段结果
     */
    public static ParentChildSplitResult ParentChildSplitter(String parentDelimiter, String childDelimiter,
                                                             int maxSegmentSize, int maxOverlapSize) {
        return new ParentChildSplitResult() {
            @Override
            public ParentChildSegments split(Document document) {

                String text = document.text();
                List<TextSegment> parentSegments = new ArrayList<>();
                List<TextSegment> childSegments = new ArrayList<>();

                // 标准化换行符
                text = text.replace("\r\n", "\n").replace("\r", "\n");

                // 如果文本不以父分隔符开头，则添加一个
                if (!text.startsWith(parentDelimiter)) {
                    text = parentDelimiter + "\n" + text;
                }

                // 按父分隔符分割
                String[] parentParts = text.split(Pattern.quote(parentDelimiter));

                for (int parentIndex = 0; parentIndex < parentParts.length; parentIndex++) {
                    String parentPart = parentParts[parentIndex].trim();
                    if (oConvertUtils.isEmpty(parentPart)) {
                        continue;
                    }

                    // 获取父分段标题（第一行非空内容）
                    String[] lines = parentPart.split("\n");
                    String parentTitle = "";
                    for (String line : lines) {
                        line = line.trim();
                        if (oConvertUtils.isNotEmpty(line)) {
                            parentTitle = line;
                            break;
                        }
                    }
                    if (oConvertUtils.isEmpty(parentTitle)) {
                        parentTitle = "parent_" + parentIndex;
                    }
                    // 生成父分段ID
                    String parentId = document.metadata().getString("source") + "::p" + parentIndex;

                    // 创建父分段
                    Metadata parentMetadata = document.metadata().copy();
                    parentMetadata.put("parent_id", parentId);
                    parentMetadata.put("parent_index", parentIndex);
                    parentMetadata.put("parent_title", parentTitle);
                    parentMetadata.put("segment_type", "parent");

                    parentSegments.add(TextSegment.from(parentPart, parentMetadata));

                    // 在父分段内按子分隔符分割
                    String childText = parentPart;
                    if (!childText.startsWith(childDelimiter)) {
                        childText = childDelimiter + "\n" + childText;
                    }

                    String[] childParts = childText.split(Pattern.quote(childDelimiter));

                    for (int childIndex = 0; childIndex < childParts.length; childIndex++) {
                        String childPart = childParts[childIndex].trim();
                        if (oConvertUtils.isEmpty(childPart)) {
                            continue;
                        }

                        // 如果子分段太长，进一步分割
                        if (childPart.length() > maxSegmentSize) {
                            List<TextSegment> subSegments = splitLongTextForChild(childPart, maxSegmentSize, maxOverlapSize,
                                    parentId, parentIndex, childIndex, parentTitle, document.metadata());
                            childSegments.addAll(subSegments);
                        } else {
                            // 创建子分段元数据
                            Metadata childMetadata = document.metadata().copy();
                            childMetadata.put("parent_id", parentId);
                            childMetadata.put("parent_index", parentIndex);
                            childMetadata.put("child_index", childIndex);
                            childMetadata.put("parent_title", parentTitle);
                            childMetadata.put("segment_type", "child");

                            childSegments.add(TextSegment.from(childPart, childMetadata));
                        }
                    }
                }

                return new ParentChildSegments(parentSegments, childSegments);
            }
        };
    }

    /**
     * 分割过长的子分段文本
     *
     * @param text 文本
     * @param maxSegmentSize 最大分段大小
     * @param maxOverlapSize 最大重叠大小
     * @param parentId 父分段ID
     * @param parentIndex 父分段索引
     * @param childIndex 子分段索引
     * @param parentTitle 父分段标题
     * @param baseMetadata 基础元数据
     * @return 文本分段列表
     */
    private static List<TextSegment> splitLongTextForChild(String text, int maxSegmentSize, int maxOverlapSize,
                                                           String parentId, int parentIndex, int childIndex,
                                                           String parentTitle, Metadata baseMetadata) {
        List<TextSegment> segments = new ArrayList<>();
        int start = 0;
        int subChildIndex = 0;

        while (start < text.length()) {
            int end = Math.min(start + maxSegmentSize, text.length());

            // 尝试在单词边界处分割
            if (end < text.length()) {
                int lastSpace = text.lastIndexOf(' ', end);
                if (lastSpace > start) {
                    end = lastSpace;
                }
            }

            String segment = text.substring(start, end).trim();
            if (oConvertUtils.isNotEmpty(segment)) {
                Metadata segmentMetadata = baseMetadata.copy();
                segmentMetadata.put("parent_id", parentId);
                segmentMetadata.put("parent_index", parentIndex);
                segmentMetadata.put("child_index", childIndex + "_" + subChildIndex);
                segmentMetadata.put("parent_title", parentTitle);
                segmentMetadata.put("segment_type", "child");

                segments.add(TextSegment.from(segment, segmentMetadata));
                subChildIndex++;
            }

            // 计算下一个分段的起始位置，考虑重叠
            start = Math.max(start + 1, end - maxOverlapSize);
        }

        return segments;
    }

    /**
     * 分割过长的文本
     *
     * @param text 文本
     * @param maxSegmentSize 最大分段大小
     * @param maxOverlapSize 最大重叠大小
     * @param baseMetadata 基础元数据
     * @return 文本分段列表
     */
    private static List<TextSegment> splitLongText(String text, int maxSegmentSize, int maxOverlapSize, Metadata baseMetadata) {
        List<TextSegment> segments = new ArrayList<>();
        int start = 0;
        int segmentIndex = 0;

        while (start < text.length()) {
            int end = Math.min(start + maxSegmentSize, text.length());

            // 尝试在单词边界处分割
            if (end < text.length()) {
                int lastSpace = text.lastIndexOf(' ', end);
                if (lastSpace > start) {
                    end = lastSpace;
                }
            }

            String segment = text.substring(start, end).trim();
            if (oConvertUtils.isNotEmpty(segment)) {
                Metadata segmentMetadata = baseMetadata.copy();
                segmentMetadata.put("index", segmentIndex++);
                segments.add(TextSegment.from(segment, segmentMetadata));
            }

            // 计算下一个分段的起始位置，考虑重叠
            start = Math.max(start + 1, end - maxOverlapSize);
        }

        return segments;
    }

    /**
     * 创建父子分段器（返回标准DocumentSplitter类型）
     * 可以直接与EmbeddingStoreIngestor配合使用
     * 父分段和子分段都在同一个List中返回，通过metadata的segment_type字段区分
     *
     * @param parentDelimiter 父分隔符，如"####"
     * @param childDelimiter 子分隔符，如"@@@"
     * @param maxSegmentSize 最大分段大小
     * @param maxOverlapSize 最大重叠大小
     * @return DocumentSplitter
     */
    public static DocumentSplitter ParentChildDocumentSplitter(String parentDelimiter, String childDelimiter,
                                                               int maxSegmentSize, int maxOverlapSize) {
        return new DocumentSplitter() {
            @Override
            public List<TextSegment> split(Document document) {
                String text = document.text();
                List<TextSegment> allSegments = new ArrayList<>();

                // 标准化换行符
                text = text.replace("\r\n", "\n").replace("\r", "\n");

                // 如果文本不以父分隔符开头，则添加一个
                if (!text.startsWith(parentDelimiter)) {
                    text = parentDelimiter + "\n" + text;
                }

                // 按父分隔符分割
                String[] parentParts = text.split(Pattern.quote(parentDelimiter));

                for (int parentIndex = 0; parentIndex < parentParts.length; parentIndex++) {
                    String parentPart = parentParts[parentIndex].trim();
                    if (oConvertUtils.isEmpty(parentPart)) {
                        continue;
                    }

                    // 获取父分段标题（第一行非空内容）
                    String[] lines = parentPart.split("\n");
                    String parentTitle = "";
                    for (String line : lines) {
                        line = line.trim();
                        if (oConvertUtils.isNotEmpty(line)) {
                            parentTitle = line;
                            break;
                        }
                    }
                    if (oConvertUtils.isEmpty(parentTitle)) {
                        parentTitle = "parent_" + parentIndex;
                    }

                    // 生成父分段ID
                    String parentId = document.metadata().getString("source") + "::p" + parentIndex;

                    // 创建父分段
                    Metadata parentMetadata = document.metadata().copy();
                    parentMetadata.put("parent_id", parentId);
                    parentMetadata.put("parent_index", parentIndex);
                    parentMetadata.put("parent_title", parentTitle);
                    parentMetadata.put("segment_type", "parent");

                    allSegments.add(TextSegment.from(parentPart, parentMetadata));

                    // 在父分段内按子分隔符分割
                    String childText = parentPart;
                    if (!childText.startsWith(childDelimiter)) {
                        childText = childDelimiter + "\n" + childText;
                    }

                    String[] childParts = childText.split(Pattern.quote(childDelimiter));

                    for (int childIndex = 0; childIndex < childParts.length; childIndex++) {
                        String childPart = childParts[childIndex].trim();
                        if (oConvertUtils.isEmpty(childPart)) {
                            continue;
                        }

                        // 如果子分段太长，进一步分割
                        if (childPart.length() > maxSegmentSize) {
                            List<TextSegment> subSegments = splitLongTextForChild(childPart, maxSegmentSize, maxOverlapSize,
                                    parentId, parentIndex, childIndex, parentTitle, document.metadata());
                            allSegments.addAll(subSegments);
                        } else {
                            // 创建子分段元数据
                            Metadata childMetadata = document.metadata().copy();
                            childMetadata.put("parent_id", parentId);
                            childMetadata.put("parent_index", parentIndex);
                            childMetadata.put("child_index", childIndex);
                            childMetadata.put("parent_title", parentTitle);
                            childMetadata.put("segment_type", "child");

                            allSegments.add(TextSegment.from(childPart, childMetadata));
                        }
                    }
                }

                return allSegments;
            }
        };
    }

    /**
     * 父子分段结果接口
     */
    public interface ParentChildSplitResult {
        ParentChildSegments split(Document document);
    }

    /**
     * 父子分段结果类
     */
    public static class ParentChildSegments {
        private final List<TextSegment> parentSegments;
        private final List<TextSegment> childSegments;

        public ParentChildSegments(List<TextSegment> parentSegments, List<TextSegment> childSegments) {
            this.parentSegments = parentSegments;
            this.childSegments = childSegments;
        }

        public List<TextSegment> getParentSegments() {
            return parentSegments;
        }

        public List<TextSegment> getChildSegments() {
            return childSegments;
        }
    }
}
